# cloudflare_workers_ai_demo.py
from langchain.agents import AgentExecutor, Tool, initialize_agent
from langchain.chains import LLMChain
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain.memory import ConversationBufferMemory
from langchain_openai import ChatOpenAI
from dotenv import load_dotenv
import os

# 加载环境变量
load_dotenv()

# ================== Cloudflare Workers AI 配置 ==================
CLOUDFLARE_API_KEY = os.getenv("CLOUDFLARE_API_KEY")
ACCOUNT_ID = os.getenv("CLOUDFLARE_ACCOUNT_ID")  # 替换为你的账户ID
MODEL_NAME = "@cf/meta/llama-2-7b-chat-fp16"     # 可更换其他Cloudflare支持的模型

# 初始化 LLM
llm = ChatOpenAI(
    openai_api_base=f"https://api.cloudflare.com/client/v4/accounts/{ACCOUNT_ID}/ai/run/",
    model=MODEL_NAME,
    temperature=0.7,
    max_tokens=256
)

# ================== 自定义 Tool 示例 ==================
from langchain.tools import BaseTool
from typing import Optional

class SimplePrintTool(BaseTool):
    name = "simple_printer"
    description = "当需要打印简单消息时使用此工具"

    def _run(self, message: str) -> str:
        """执行打印操作的逻辑"""
        print(f"[工具输出] {message}")
        return f"已成功打印：{message}"

# ================== 创建 Chain ==================
prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一个智能助手，可以调用工具处理用户请求"),
    MessagesPlaceholder(variable_name="chat_history"),
    ("human", "{input}"),
])

memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
chain = LLMChain(llm=llm, prompt=prompt, memory=memory)

# ================== 初始化 Agent ==================
tools = [
    SimplePrintTool(),
    # 可在此添加更多工具...
]

agent = initialize_agent(
    tools,
    llm,
    agent="chat-conversational-react-description",
    verbose=True,
    memory=memory,
    handle_parsing_errors=True
)

# ================== 执行示例 ==================
if __name__ == "__main__":
    # 测试 Chain
    print("\n=== 测试 Chain ===")
    chain_response = chain.invoke({"input": "写一句关于Cloudflare的诗"})
    print(chain_response["text"])

    # 测试 Agent 使用工具
    print("\n=== 测试 Agent 使用工具 ===")
    agent_response = agent.invoke({
        "input": "请用 simple_printer 工具打印这句话：Cloudflare Workers AI 很棒！"
    })
    print(agent_response["output"])